# Copyright (c) 2021 Presto Labs Pte. Ltd.
# Author: jingyuan

import json
import os
import shutil
import subprocess
import tempfile
import logging
import time
import bisect

import pandas as pd
from coin.exchange.binance_futures_swap.kr_rest.native_private_client import \
    BinanceSwapNativePrivateClient
from coin.exchange.binance_futures.kr_rest.native_private_client import \
    BinanceFuturesNativePrivateClient
from coin2.support.product_info.product_info_holder import ProductInfoHolder
from cc.appcoin2.strategy.support.constants import LEVERGAE_MAP_LIST
from cc.appcoin2.strategy.support.util import get_cmd_vault_key_copy
from python.coin_deploy.basis_strat2.script.adjust_universe import get_max_pos

from absl import app, flags


class MyCoinSwapClient(BinanceFuturesNativePrivateClient):
  def __init__(self, key_file):
    super().__init__(key_file=key_file)
    ph = ProductInfoHolder()
    pi_df = ph.df_pi[['mea', 'native_symbol', 'contract_value']]
    pi_df = pi_df[pi_df['mea'] == 'Futures.Binance.v1-delivery'][['native_symbol', 'contract_value']]
    self.pi_df = pi_df

  def set_leverage(self, symbol, leverage):
    method = 'POST'
    path = '/dapi/v1/leverage'
    params = {
        'symbol': symbol,
        'leverage': leverage,
        'recvWindow': self._recv_window,
        'timestamp': int(time.time() * 1000),
    }
    update = self._query(method, path, params)
    return update

  def _get_leverage(self, symbol):
    method = 'GET'
    path = '/dapi/v1/positionRisk'
    params = {
        'symbol': symbol,
        'recvWindow': self._recv_window,
        'timestamp': int(time.time() * 1000),
    }
    update = self._query(method, path, params)
    return update

  def get_leverage(self, filter_zero=True):
    results = self._get_leverage('')
    df = pd.DataFrame(results.as_json()['msg'],
                      columns=[
                          'symbol', 'leverage', 'maxQty', 'positionAmt', 'entryPrice',
                          'markPrice', 'unRealizedProfit'
                      ])
    df = pd.merge(df, self.pi_df, how='left', left_on='symbol', right_on='native_symbol')
    df.set_index('symbol', inplace=True)
    df = df.apply(pd.to_numeric, errors='ignore')
    df['positionUSD'] = (df.positionAmt * df.contract_value).round(2)
    df['maxNotionalValue'] = (df.maxQty * df.markPrice).round(2)
    df = df[[
        'leverage', 'maxNotionalValue', 'positionUSD', 'positionAmt', 'entryPrice', 'markPrice',
        'unRealizedProfit'
    ]]
    lever_info = df['leverage'].to_dict()
    if filter_zero:
      return df[df.positionAmt != 0], lever_info
    return df, lever_info


class MySwapClient(BinanceSwapNativePrivateClient):
  def __init__(self, key_file):
    super().__init__(key_file=key_file)
    self.leverage_df = None
    self.brackets = None

  def set_leverage(self, symbol, leverage):
    method = 'POST'
    path = '/fapi/v1/leverage'
    params = {
        'symbol': symbol,
        'leverage': leverage,
        'recvWindow': self._recv_window,
        'timestamp': int(time.time() * 1000),
    }
    update = self._query(method, path, params)
    return update

  def _get_leverage(self, symbol):
    method = 'GET'
    path = '/fapi/v2/positionRisk'
    params = {
        'symbol': symbol,
        'recvWindow': self._recv_window,
        'timestamp': int(time.time() * 1000),
    }
    update = self._query(method, path, params)
    return update

  def get_leverage(self, filter_zero=True):
    results = self._get_leverage('')
    df = pd.DataFrame(results.as_json()['msg'],
                      columns=[
                          'symbol', 'leverage', 'maxNotionalValue', 'positionAmt', 'entryPrice',
                          'markPrice', 'unRealizedProfit'
                      ])
    df.set_index('symbol', inplace=True)
    df = df.apply(pd.to_numeric, errors='ignore')
    df['positionUSD'] = (df.positionAmt * df.markPrice).round(2)
    df = df[[
        'leverage', 'maxNotionalValue', 'positionUSD', 'positionAmt', 'entryPrice', 'markPrice',
        'unRealizedProfit'
    ]]
    lever_info = df['leverage'].to_dict()
    if filter_zero:
      return df[df.positionAmt != 0], lever_info
    return df, lever_info

  def _get_leverage_brackets(self):
    method = 'GET'
    path = '/fapi/v1/leverageBracket'
    params = {
        'symbol': '',
        'recvWindow': self._recv_window,
        'timestamp': int(time.time() * 1000),
    }
    update = self._query(method, path, params)
    return update

  def get_leverage_brackets(self):
    results = self._get_leverage_brackets()
    dfs = []
    brackets = {}
    for msg in results.as_json()['msg']:
      df = pd.DataFrame(msg['brackets'],
                        columns=[
                          'bracket', 'initialLeverage', 'notionalCap', 'notionalFloor', 'maintMarginRatio', 'cum'
                      ])
      df = df.apply(pd.to_numeric, errors='ignore')
      brackets[msg['symbol']] = sorted(set(df['notionalCap'].values))
      df['symbol'] = msg['symbol']
      dfs.append(df)
    df = pd.concat(dfs)
    self.leverage_df = df
    self.brackets = brackets

  def get_suitable_leverage(self, symbol, notional_cap, leverage):
    if not self.brackets and not self.leverage_df:
      self.get_leverage_brackets()
    if symbol not in self.brackets.keys() or len(self.brackets[symbol]) == 0:
      logging.warning(f"{symbol} does not have/ have empty leverage brackets.")
      return leverage
    idx = bisect.bisect(self.brackets[symbol], notional_cap - 1)
    if idx == len(self.brackets[symbol]):
      idx = idx - 1
      logging.warning(
          f'Assigned notional cap {notional_cap} exceed largest bracket {self.brackets[symbol][-1]}')
    leverage = self.leverage_df.loc[(self.leverage_df['symbol'] == symbol) & (
        self.leverage_df['notionalCap'] == self.brackets[symbol][idx]), 'initialLeverage'].values[0]
    return leverage


def init_flags():
  flags.DEFINE_string('strategy', '', 'strategy_name')
  flags.DEFINE_string('account', '', 'account name')
  flags.DEFINE_string('key_filepath', '', 'key filepath')
  flags.DEFINE_string('symbol', 'ZILUSDT', 'native symbol')
  flags.DEFINE_integer('leverage', 10, 'upd leverage')
  flags.DEFINE_integer('notional_cap', None, 'notional cap')
  flags.DEFINE_bool('query_only', True, 'check leverage')
  flags.DEFINE_bool(
      'filter_zero',
      True,
      'filter symbols has no positions and leverage=20 in display')
  flags.DEFINE_bool('update_all', False, 'update all symbols')
  flags.DEFINE_bool('is_coin_swap', False, 'for coin margin swap')
  flags.DEFINE_bool('is_fund_account', False, 'for fund account')


def get_binance_account(strategy):
  if os.path.exists(f'python/coin_deploy/basis_strat2/driver/{strategy}.json'):
    driver = json.load(open(f'python/coin_deploy/basis_strat2/driver/{strategy}.json'))
    for exch, config in driver['order'].items():
      if config['mea'] == 'Futures.Binance.v1':
        return config['connection_config']['key_filepath'].split('/')[3]
  raise ValueError(f"look up account for {strategy} failed")


def check_max_pos(df, max_pos):
  products = []
  config_path = f'python/coin_deploy/basis_strat2/config/{flags.FLAGS.strategy}.json'
  if not os.path.exists(config_path):
    return
  with open(config_path) as fi:
    config = json.load(fi)
  for key in config['products']:
    if config['products'][key].get('trade_symbol', False):
      products.append(config['products'][key]['trade_symbol'])
    elif config['products'][key].get('symbol', False):
      products.append(config['products'][key]['symbol'])
  products = [symbol.split('.')[0].split(':')[-1].replace('-', '') for symbol in products]

  for row in df.iterrows():
    symbol = row[0]
    if symbol not in products:
      continue
    max_notional_value = float(row[1]['maxNotionalValue'])
    if max_pos > max_notional_value:
      logging.warning(
          f'{flags.FLAGS.strategy} {symbol} max_allowed_postion {max_notional_value} is smaller than the current max position {max_pos}')


def main(_):
  pd.set_option('display.max_rows', None)
  if flags.FLAGS.strategy and not flags.FLAGS.account:
    account = get_binance_account(flags.FLAGS.strategy)
    print(f'{flags.FLAGS.strategy} is using account {account}')
  else:
    account = flags.FLAGS.account
  tmp_dir = tempfile.mkdtemp()
  try:
    if flags.FLAGS.key_filepath:
      key_file = flags.FLAGS.key_filepath
    else:
      key_file = os.path.join(tmp_dir, 'key.json')
      vault_key_name = f"trade/{account}/binance/fut_trade_key.json"
      cmd = get_cmd_vault_key_copy(vault_key_name, key_file, flags.FLAGS.is_fund_account)
      output = subprocess.check_output(cmd, shell=True)
    if flags.FLAGS.is_coin_swap:
      msc = MyCoinSwapClient(key_file=key_file)
    else:
      msc = MySwapClient(key_file=key_file)

    results, lever_info = msc.get_leverage(flags.FLAGS.filter_zero)
    print("\n#####################")
    print("-- Before Adjust --")
    print(results)
    print("######################\n")
    if not flags.FLAGS.query_only:
      if flags.FLAGS.update_all:
        symbols = lever_info.keys()
        flags.FLAGS.symbol = ",".join(symbols)
      assert flags.FLAGS.symbol
      for s in flags.FLAGS.symbol.split(','):
        for key in LEVERGAE_MAP_LIST:
          s = s.replace(key, LEVERGAE_MAP_LIST[key])
        if "USD" not in s:
          s += "USDT"
        if flags.FLAGS.notional_cap:
          leverage = msc.get_suitable_leverage(s, flags.FLAGS.notional_cap, flags.FLAGS.leverage)
          leverage = min(leverage, 10)
        else:
          leverage = flags.FLAGS.leverage
        if s in lever_info and leverage == lever_info[s]:
          print("Already applied. skip.. ", s, ", leverage =", lever_info[s])
          continue
        result = msc.set_leverage(s, leverage)
        print(result)
      print("\n#####################")
      print("-- After Adjust --")
      res = msc.get_leverage(flags.FLAGS.filter_zero)[0]
      print(res)
      print("######################\n")
      if not flags.FLAGS.notional_cap:
        max_pos = get_max_pos(flags.FLAGS.strategy)
        if max_pos:
          check_max_pos(res, max_pos)
      else:
        check_max_pos(res, flags.FLAGS.notional_cap)
  except BaseException:
    raise
  finally:
    shutil.rmtree(tmp_dir)


if __name__ == '__main__':
  init_flags()
  app.run(main)

# Usage
# Check leverage
# ./pyrunner cc/appcoin2/strategy/support/binance_leverage.py --account=xunke12_swap --symbol=
# Update leverage
# ./pyrunner cc/appcoin2/strategy/support/binance_leverage.py --account=xunke12_swap --symbol=ENJUSDT --leverage=19 --query_only=False
# ./pyrunner cc/appcoin2/strategy/support/binance_leverage.py --account=xunke51_swap  --leverage=10 --query_only=0 --symbol VETUSDT,BTTUSDT,KNCUSDT,ENJUSDT,XEMUSDT
# update by strategy
# ./pyrunner cc/appcoin2/strategy/support/binance_leverage.py --strategy=basis_smm2_binance_spot_binance_perpetual  --query_only=0 --symbol=ICXUSDT --leverage=5
# ./pyrunner cc/appcoin2/strategy/support/binance_leverage.py --strategy=basis_smm2_binance_spot_binance_perpetual  --query_only=0 --symbol=SXP,DOGEBUSD,BTCUSDT_210924 --leverage=5
# coin margin swap
# ./pyrunner cc/appcoin2/strategy/support/binance_leverage.py --is_coin_swap --account=xunkeXX --leverage=5 --query_only=0 --symbol=SOLUSD_PERP
# use key_filepath
# ./pyrunner cc/appcoin2/strategy/support/binance_leverage.py --leverage=5 --query_only=0 --symbol=1000LUNCBUSD --key_filepath ~/tmp/tmp8rwcuof3/key.json
# ./pyrunner cc/appcoin2/strategy/support/binance_leverage.py --leverage=5 --query_only=0 --update_all --key_filepath ~/tmp/tmp8rwcuof3/key.json
